// Copyright 2022 Huawei Cloud Computing Technology Co., Ltd.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "CaeTlsSocket.h"
#include "../cae_common/CaeLogConfig.h"
#include "../cae_agent/CaeAgent.h"
#include <cerrno>
#include <cstring>
#include <cstdio>
#include <cstdlib>
#include <sys/socket.h>
#include <linux/in.h>
#include <linux/tcp.h>
#include <unistd.h>

using namespace std;

CaeTlsSocket::CaeTlsSocket()
{
    m_fd = -1;
    m_status = -1;
    m_localIp = 0;
    m_remoteIp = 0;
    m_localPort = 0;
    m_remotePort = 0;
    m_socketOption = 0;
    m_eventNotice = nullptr;
    m_ctx = nullptr;
    m_ssl = nullptr;
    pthread_mutex_init(&m_socketMutex, nullptr);
    std::string caCertPath = CaeConfigManage::GetCaCertPath();
    std::string certPath = CaeConfigManage::GetCertPath();
    std::string privKeyPath = CaeConfigManage::GetPrivKeyPath();
    std::string certUserData = CaeConfigManage::GetCertUserData();
    if (caCertPath.length() == 0) {
        LOGE("Get ca_crt failed.");
        return;
    }
    if (certPath.length() == 0) {
        LOGE("Get cae_crt failed.");
        return;
    }
    if (privKeyPath.length() == 0) {
        LOGE("Get cae_key failed.");
        return;
    }

    if (SSL_library_init() != 1) {
        LOGE("Init ssl library failed: %s.", ERR_reason_error_string(ERR_get_error()));
        return;
    }
    OpenSSL_add_all_algorithms();
    SSL_load_error_strings();

    m_ctx = SSL_CTX_new(TLS_server_method());
    if (m_ctx == nullptr) {
        LOGE("Create SSL_CTX failed: %s.", ERR_reason_error_string(ERR_get_error()));
    }
    SSL_CTX_set_min_proto_version(m_ctx, TLS1_2_VERSION);
    SSL_CTX_set_verify(m_ctx, SSL_VERIFY_NONE, nullptr);
    if (!certUserData.empty()) {
        LOGI("Set cert user data.");
        SSL_CTX_set_default_passwd_cb_userdata(m_ctx, (void *)certUserData.c_str());
    }

    if (SSL_CTX_use_certificate_file(m_ctx, certPath.c_str(), SSL_FILETYPE_PEM) != 1) {
        LOGE("Server cert load failed: %s.", ERR_reason_error_string(ERR_get_error()));
        SSL_CTX_free(m_ctx);
        m_ctx = nullptr;
        return;
    }

    if (SSL_CTX_use_PrivateKey_file(m_ctx, privKeyPath.c_str(), SSL_FILETYPE_PEM) != 1) {
        LOGE("Server private key load failed: %s.", ERR_reason_error_string(ERR_get_error()));
        SSL_CTX_free(m_ctx);
        m_ctx = nullptr;
        return;
    }

    if (SSL_CTX_check_private_key(m_ctx) != 1) {
        LOGE("Private key does not match the certificate public key.");
        SSL_CTX_free(m_ctx);
        m_ctx = nullptr;
        return;
    }

    const char allowedCiphers[] = "ECDHE-RSA-AES256-GCM-SHA384:ECDHE-ECDSA-AES256-GCM-SHA384";
    if (SSL_CTX_set_cipher_list(m_ctx, allowedCiphers) != 1) {
        LOGE("SSL_CTX_set_cipher_list failed: %s.", ERR_reason_error_string(ERR_get_error()));
    }
    SSL_CTX_set_mode(m_ctx, SSL_MODE_AUTO_RETRY);

    m_ssl = SSL_new(m_ctx);
    if (m_ssl == nullptr) {
        LOGE("Failed to create SSL: %s.", ERR_reason_error_string(ERR_get_error()));
        SSL_CTX_free(m_ctx);
        m_ctx = nullptr;
        return;
    }
}

CaeTlsSocket::~CaeTlsSocket()
{
    m_status = -1;
    if (-1 != m_fd) {
        shutdown(m_fd, SHUT_RDWR);
        close(m_fd);
    }
    m_fd = -1;
    if (m_ctx != nullptr) {
        SSL_CTX_free(m_ctx);
        m_ctx = nullptr;
    }
    if (m_ssl != nullptr) {
        SSL_free(m_ssl);
        m_ssl = nullptr;
    }
}

void CaeTlsSocket::CloseSocket()
{
    m_status = -1;
    if (m_fd != -1) {
        shutdown(m_fd, SHUT_RDWR);
        close(m_fd);
    }
    m_fd = -1;
}

bool CaeTlsSocket::ConfigSSL()
{
    if (m_ctx == nullptr) {
        LOGE("Failed to config ssl, m_ctx is NULL.");
        return false;
    }
    if (nullptr == m_ssl) {
        LOGE("Failed to config ssl, m_ssl is NULL.");
        SSL_CTX_free(m_ctx);
        m_ctx = nullptr;
        return false;
    }

    if (1 != SSL_set_fd(m_ssl, m_fd)) {
        LOGE("Failed to config ssl, set fd failed: %s.", ERR_reason_error_string(ERR_get_error()));
        return false;
    }

    fd_set wr_handle;
    fd_set rd_handle;
    struct timeval tv{};
    int sslErr;
    int sslRet;
    bool bRetry = true;
    int iTryTimes = 0;
    const int retryMaxTime = 4;
    do {
        FD_ZERO(&wr_handle);
        FD_ZERO(&rd_handle);
        sslRet = ::SSL_accept(m_ssl);
        if (sslRet == 1) {
            break;
        }
        sslErr = ::SSL_get_error(m_ssl, sslRet);
        switch (sslErr) {
            case SSL_ERROR_WANT_WRITE:
                FD_SET(m_fd, &wr_handle);
                bRetry = true;
                LOGI("SSL get error: SSL_WANT_WRITE.");
                break;
            case SSL_ERROR_WANT_READ:
                FD_SET(m_fd, &rd_handle);
                bRetry = true;
                LOGI("SSL get error: SSL_WANT_READ.");
                break;
            case SSL_ERROR_ZERO_RETURN:
                bRetry = false;
                LOGE("SSL get error: SSL_ERROR_ZERO_RETURN.");
                break;
            case SSL_ERROR_SYSCALL:
                LOGI("SSL get error: SSL_ERROR_SYSCALL.");
                bRetry = true;
                if (SSL_want_read(m_ssl)) {
                    FD_SET(m_fd, &rd_handle);
                } else if (SSL_want_write(m_ssl)) {
                    FD_SET(m_fd, &wr_handle);
                } else {
                    LOGE("SSL get error: SSL_ERROR_SYSCALL, An I/O error has occurred error %d, ", errno);
                    bRetry = false;
                }
                break;
            default:
                LOGE("SSL get error: [%d]", sslErr);
                bRetry = false;
                break;
        }
        if (bRetry) {
            tv.tv_sec = 0;
            tv.tv_usec = 500000;
            int result = select(m_fd + 1, &rd_handle, &wr_handle, nullptr, &tv);
            if ((result < 0) && (514 != errno) && (EINTR != errno)) {
                LOGE("Failed to select SSL_getError(%d), try times(%d).", sslErr, iTryTimes);
                bRetry = false;
            }
        }
        iTryTimes++;
        if (iTryTimes >= retryMaxTime) {
            LOGD("SSL_accept() failed, too many attempts(%d).", iTryTimes);
            bRetry = false;
        }
        LOGD("Try times is %d", iTryTimes);
    } while (bRetry);

    if (sslRet != 1) {
        LOGE("Failed to get ssl accept, result is %d.", sslRet);
        return false;
    }

    if (SSL_get_verify_result(m_ssl) != X509_V_OK) {
        LOGE("Failed to get ssl verify, result is %s.", ERR_reason_error_string(ERR_get_error()));
        return false;
    }
    return true;
}

int CaeTlsSocket::Send(void *pPkt, size_t size)
{
    int byte;
    int pktSize = static_cast<int>(size);
    int err = pktSize;
    if (nullptr == pPkt) {
        return SOCKET_SEND_FAIL_RETRY;
    }
    pthread_mutex_lock(&m_socketMutex);
    signal(SIGPIPE, SIG_IGN);
    while (true) {
        byte = SSL_write(m_ssl, pPkt, pktSize);
        if (pktSize == byte) {
            pthread_mutex_unlock(&m_socketMutex);
            return pktSize;
        } else if ((0 <= byte) && (pktSize > byte)) {
            pktSize -= byte;
            pPkt = (void *)((uint8_t *)pPkt + byte);
            continue;
        } else if (!((EAGAIN == errno) || (EWOULDBLOCK == errno) || (EINTR == errno))) {
            err = SOCKET_SEND_FAIL_DISCONN;
            LOGE("Socket Send error, errno = %d.", errno);
            m_status = SOCKET_STATUS_ERROR;
            break;
        } else if (m_fd == -1) {
            err = SOCKET_SEND_FAIL_DISCONN;
            LOGE("Socket Send error(fd = -1), errno = %d.", errno);
            CaeAgent::GetInstance().CloseClient();
            break;
        } else {
            LOGD("Socket Send error, errno=%d", errno);
            break;
        }
    }
    pthread_mutex_unlock(&m_socketMutex);
    this->EventNotice();
    return err;
}


int CaeTlsSocket::Recv(void *pPkt, size_t size)
{
    uint32_t flag = m_socketOption & SOCKET_OPTION_BITSET_QUICKACK;
    int byte;
    int socketFd = this->m_fd;
    if (nullptr == pPkt) {
        return SOCKET_RECV_FAIL_RETRY;
    }
    byte = SSL_read(m_ssl, pPkt, static_cast<int>(size));
    if (byte <= 0) {
        flag = 0;
        if (!((byte < 0) && ((EAGAIN == errno) || (EWOULDBLOCK == errno) || (EINTR == errno)))) {
            m_status = SOCKET_STATUS_ERROR;
            LOGE("Socket receive error, #errno = %d, #errInfo is %s.", errno, ERR_reason_error_string(ERR_get_error()));
            byte = SOCKET_RECV_FAIL_DISCONN;
        }
    }
    if (flag) {
        setsockopt(socketFd, IPPROTO_TCP, TCP_QUICKACK, (int[]) { 1 }, sizeof(int));
    }

    return byte;
}

void CaeTlsSocket::SetSocketTimeout(long sec, long uSec)
{
    struct timeval timeout{};
    timeout.tv_sec = sec;
    timeout.tv_usec = uSec;

    if (setsockopt(m_fd, SOL_SOCKET, SO_SNDTIMEO, &timeout, sizeof(timeout)) < 0) {
        LOGD("Set socket opt SO_SNDTIMEO failed, errno: (%d): %s.", errno, strerror(errno));
    }

    if (setsockopt(m_fd, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout)) < 0) {
        LOGD("Set socket opt SO_SNDTIMEO failed, errno: (%d): %s.", errno, strerror(errno));
    }
}
